import { enumStrokeColor } from '../theme'

const getPolylineModel = PolylineEdgeModel => {
  class PolylineModel extends PolylineEdgeModel {
    constructor(data, graphModel) {
      super(data, graphModel)
      const { type } = data.properties || {}
      const pStroke = enumStrokeColor[type] || enumStrokeColor.primary
      this.stroke = pStroke
    }
    // 在生命周期中，没有执行该方法
    setAttributes() {
      const {
        stroke,
        properties: { type },
      } = this
      const pStroke = enumStrokeColor[type] || stroke
      this.stroke = pStroke
    }
  }
  return PolylineModel
}

export const registerPolylineEdge = ({ PolylineEdge, PolylineEdgeModel }) => ({
  model: getPolylineModel(PolylineEdgeModel),
  view: PolylineEdge,
})
